# _*_ coding : utf-8 _*_
# @Time : 2025/5/2 16:53
# @Author : liuchengjie
# @File : WS2812_LED
# @Project : esp32
import machine
import time


class WS2812_LED:

    def __init__(self, pin, led_count, brightness=1.0):
        self.pin = pin
        self.led_count = led_count
        self.brightness = brightness
        self.buffer = bytearray(led_count * 3)  # 每个LED需要3个字节(RGB)

    def set_pixel(self, index, r, g, b):
        if 0 <= index < self.led_count:
            # 应用亮度
            r = int(r * self.brightness)
            g = int(g * self.brightness)
            b = int(b * self.brightness)

            # WS2812使用GRB顺序
            self.buffer[index * 3] = g
            self.buffer[index * 3 + 1] = r
            self.buffer[index * 3 + 2] = b

    def fill(self, r, g, b):
        for i in range(self.led_count):
            self.set_pixel(i, r, g, b)

    def show(self):
        try:
            self.pin.drive(machine.Pin.OPEN_DRAIN)
            timing = (400, 850, 800, 450)  # 高电平0, 低电平0, 高电平1, 低电平1 (ns)
            machine.bitstream(self.pin, 0, timing, self.buffer)
        except AttributeError:
            try:
                import neopixel
                np = neopixel.NeoPixel(self.pin, self.led_count)
                for i in range(self.led_count):
                    np[i] = (self.buffer[i * 3 + 1], self.buffer[i * 3], self.buffer[i * 3 + 2])
                np.write()

                # 如果两种方法都不可用,使用手动位打点实现
            except ImportError:
                self._write_manual()

    def _write_manual(self):
        # 禁用中断以确保精确时序
        irq_state = machine.disable_irq()
        for byte in self.buffer:
            for bit in range(8):
                if byte & (0x80 >> bit):
                    # 发送1代码
                    self.pin.value(1)
                    time.sleep_us(1)  # ~800ns高
                    self.pin.value(0)
                    time.sleep_us(1)  # ~450ns低
                else:
                    # 发送0代码
                    self.pin.value(1)
                    time.sleep_us(1)  # ~400ns高
                    self.pin.value(0)
                    time.sleep_us(1)  # ~850ns低

            # 恢复中断
        machine.enable_irq(irq_state)

        # 重置信号, >50us的低电平作为帧结束标志
        time.sleep_us(60)

    def set_brightness(self, brightness):
        self.brightness = max(0.0, min(1.0, brightness))

    def wheel(self, pos):
        pos = pos % 256
        if pos < 85:
            return (pos * 3, 255 - pos * 3, 0)
        elif pos < 170:
            pos -= 85
            return (255 - pos * 3, 0, pos * 3)
        else:
            pos -= 170
            return (0, pos * 3, 255 - pos * 3)

    def rainbow_cycle(self, wait_ms=20):
        for j in range(256):
            for i in range(self.led_count):
                pos = (i * 256 // self.led_count + j) % 256
                r, g, b = self.wheel(pos)
                self.set_pixel(i, r, g, b)
            self.show()
            time.sleep_ms(wait_ms)
